import linecache
import os
import random
import shutil
import time

import cv2
import numpy as np


class Config:
    root = r'D:\PythonProject\face48\full'
    target = r'D:\PythonProject\face48\my'
    test = r'D:\PythonProject\face48\test'
    txt_root = r'D:\PythonProject\face48\train.txt'
    val_txt = r'D:\PythonProject\face48\val.txt'
    train_txt = r'D:\PythonProject\face48\trn.txt'
    train_batch_size = 32
    train_number_epochs = 30


def cv_imread(file_path=""):
    # file_path_gbk = file_path.encode('gbk')        # unicode转gbk，字符串变为字节数组
    # img_mat = cv2.imread(file_path_gbk.decode())  # 字节数组直接转字符串，不解码
    cv_img = cv2.imdecode(np.fromfile(file_path, dtype=np.uint8), -1)
    return cv_img


def readimage(rootpath, savepath):
    image_dir = sorted(os.listdir(rootpath))
    # 对list中图片逐一进行检查,找出其中的人脸然后写到目标文件夹下
    label = 0
    face_cascade = cv2.CascadeClassifier(r'D:\anaconda\envs\ml\Lib\site-packages\cv2\data\haarcascade_frontalface_alt.xml')
    for name in image_dir:
        if label < 12:
            label += 1
            continue
        # 创建保存路径
        withlabel = False
        if withlabel:
            targetPath = os.path.join(savepath, 's' + str(label))
        else:
            targetPath = savepath
        if not os.path.exists(targetPath):
            os.makedirs(targetPath)
        ppath = os.path.join(rootpath, name)
        # 识别人脸
        ImagePaths = os.listdir(ppath)
        count = 1
        for imagePath in ImagePaths:
            if 'jpg' not in imagePath:
                continue
            img = cv_imread(os.path.join(ppath, imagePath))
            facesize = 150
            faces = face_cascade.detectMultiScale(img, 1.1, 6, minSize=(facesize, facesize))

            if len(faces) != 0:
                for (x, y, w, h) in faces:
                    padding = 10
                    if w >= facesize and h >= facesize:
                        # 以时间戳和读取的排序作为文件名称
                        listStr = [str(int(time.time())), str(count)]
                        count += 1
                        fileName = ''.join(listStr)
                        fileName = str(label) + '_' + fileName
                        # 扩大图片，可根据坐标调整
                        X = max(int(x) - padding, 0)
                        W = min(int(x + w + padding), img.shape[1])
                        Y = max(int(y) - padding, 0)
                        H = min(int(y + h + padding), img.shape[0])

                        f = cv2.resize(img[Y:H, X:W], (W - X, H - Y))
                        savePath = os.path.join(targetPath, f'{fileName}.jpg')
                        cv2.imwrite(savePath, f)
        print(name, ': ', label)
        label += 1

    # x = np.zeros((len(image_dir), 128, 128, 3), dtype=np.uint8)
    # y = np.zeros((len(image_dir)), dtype=np.uint8)
    # for i, file in enumerate(image_dir):
    #     x[i, :, :] = cv2.imread(os.path_.join(path_, file))
    #     # x[i, :, :] = cv2.resize(img,(128, 128)) 把resize好的保存到pre文件夹里了，不用每次处理
    #     if label:
    #         y[i] = int(file.split("_")[0])


def convert(train=True):
    if train:
        f = open(Config.txt_root, 'w')
        image_dir = sorted(os.listdir(Config.target))
        print(image_dir)

        for fname in image_dir:
            t = fname.split('_')[0]
            f.write(os.path.join(Config.target, fname) + ' ' + str(t) + '\n')

        f.close()


def make_test_image():
    face_cascade = cv2.CascadeClassifier(r'D:\anaconda\envs\ml\Lib\site-packages\cv2\data\haarcascade_frontalface_alt.xml')
    count = 0
    ImagePaths = sorted(os.listdir(Config.test))
    for imagePath in ImagePaths:
        img = cv_imread(os.path.join(Config.test, imagePath))
        faces = face_cascade.detectMultiScale(img, 1.1, 5)
        if len(faces) != 0:
            for (x, y, w, h) in faces:
                # 设置人脸宽度大于xx像素，去除较小的人脸
                facesize = 40
                if w >= facesize and h >= facesize:
                    # 以时间戳和读取的排序作为文件名称
                    listStr = [str(int(time.time())), str(count)]
                    count += 1
                    fileName = ''.join(listStr)
                    # 扩大图片，可根据坐标调整
                    padding = 30
                    X = max(int(x) - padding, 0)
                    W = min(int(x + w + padding), img.shape[1])
                    Y = max(int(y) - padding, 0)
                    H = min(int(y + h + padding), img.shape[0])

                    f = cv2.resize(img[Y:H, X:W], (W - X, H - Y))
                    # 再resize成facesize的大小
                    f = cv2.resize(f, (facesize + padding, facesize + padding))
                    savePath = os.path.join(Config.test, f'{fileName}.jpg')
                    cv2.imwrite(savePath, f)
                    print(savePath)


def make_val_text():
    ori = open(Config.txt_root, 'r')

    lines = ori.readlines()
    lastlabel = "0"
    linelist = []

    tri = open(Config.train_txt, 'w')
    valfile = open(Config.val_txt, 'w')

    for line in lines:
        line.strip('\n')
        label = line.split()[1]
        if label != lastlabel:  # 每个label随机出2个测试用的图
            index = random.randint(0, int(len(linelist) / 2) - 1)
            index1 = random.randint(int(len(linelist) / 2), len(linelist) - 1)
            valfile.write(linelist[index])
            valfile.write(linelist[index1])
            for ii in range(0, len(linelist)):
                if ii != index and ii != index1:
                    tri.write(linelist[ii])
            # print(lastlabel, index, index1, len(linelist),linelist[index],linelist[index1])
            linelist.clear()
            lastlabel = label
        linelist.append(line)

    # 最后一个label的处理
    index = random.randint(0, int(len(linelist) / 2) - 1)
    index1 = random.randint(int(len(linelist) / 2), len(linelist) - 1)
    valfile.write(linelist[index])
    valfile.write(linelist[index1])
    for ii in range(0, len(linelist)):
        if ii != index and ii != index1:
            tri.write(linelist[ii])
    print(lastlabel, index, index1, len(linelist))

    tri.close()
    valfile.close()


def split_data(rootpath):
    """rootpath中保存抓出来的头像，将其分为train和val两个文件夹"""
    image_dir = os.listdir(rootpath)

    train = os.path.join(rootpath, 'train')
    test = os.path.join(rootpath, 'val')

    for memberfolder in image_dir:
        tf = os.path.join(train, memberfolder)
        testf = os.path.join(test, memberfolder)
        os.makedirs(tf, exist_ok=True)
        os.makedirs(testf, exist_ok=True)

        curpath = os.path.join(rootpath, memberfolder)
        files = sorted(os.listdir(curpath))
        for filename in files:
            # 参数 旧路径：新路径
            p = tf if random.randint(0, 10) < 7 else testf
            shutil.move(os.path.join(curpath, filename), os.path.join(p, filename))


def copyfolder():
    rootpath = r'H:\mlface'
    source_path = os.path.abspath(r'E:\Projects\source_dir')
    target_path = os.path.abspath(r'H:\faceset')

    image_dir = os.listdir(rootpath)

    for memberfolder in image_dir:
        try:
            src = os.path.join(rootpath, memberfolder)
            src = os.path.join(src, 'capture_image')
            tar = os.path.join(target_path, memberfolder)
            # os.makedirs(tar, exist_ok=True)
            # if os.path.exists(tar):
            #     # 如果目标路径存在原文件夹的话就先删除
            #     shutil.rmtree(tar)
            shutil.copytree(src, tar)
        except Exception as e:
            continue

    print('copy dir finished!')


if __name__ == '__main__':
    # readimage(r'G:\ml2021\images\face', r'G:\ml2021\images')
    split_data(r'H:\faceset')
    # convert()
    # make_val_text()
